"""
predict.py
"""
import numpy as np
from PIL import Image
import cv2

from executor import detect_model
import torch


def get_available_device():
    # 检查torch是否有CUDA支持，即是否能用GPU
    print(torch.cuda.is_available())
    if not torch.cuda.is_available():
        return 'cpu'


    # 如果CUDA可用，它还会打印出当前默认的CUDA设备（通常是第一个GPU）
    gpu_idx = torch.cuda.current_device()
    return f'cuda:{gpu_idx}'


def predict_single_img_file_old(fn, conf_threshold=0.7, plot_result=False):
    """
    根据文件名进行单张图片推理分割
    :param plot_result: 是否输出推理图片，默认False
    :param fn: 文件名
    :param conf_threshold: 置信度
    :return: 见 predict_single_img
    """
    image = Image.open(fn).convert("RGB")
    image_np = np.array(image)
    return predict_single_img(image_np, conf_threshold, plot_result)


def get_model_and_device():
    """
    获取模型和设备的公共函数
    """
    model = detect_model.get_model(detect_model='pms-slider-detect-20250518.pt')
    device = get_available_device()
    print(f'device:{device}')
    return model, device


def process_detections(results, image, conf_threshold, plot_result=False):
    """
    处理检测结果的公共函数
    """
    detectVals = []
    # 类别名称映射
    class_names = {0: "slider"}

    # 处理推理结果
    for result in results:
        boxes = result.boxes
        if boxes is not None and boxes.shape[0] > 0:
            for index, box in enumerate(boxes):
                cls = int(result.boxes[index].cls)
                conf = result.boxes[index].conf.tolist()[0]
                if conf >= conf_threshold:
                    # 获取边界框坐标
                    box_coords = box.xyxy[0].tolist()  # [x1, y1, x2, y2]
                    # 将检测结果添加到detectVals
                    detectVals.append({
                        'cls': class_names[cls],
                        'conf': conf,
                        'box': box_coords
                    })

                    # 如果需要绘制结果
                    if plot_result:
                        # 绘制边界框
                        cv2.rectangle(image,
                                    (int(box_coords[0]), int(box_coords[1])),
                                    (int(box_coords[2]), int(box_coords[3])),
                                    (0, 255, 0), 2)

                        # 准备标题文本
                        title = f"{class_names[cls]}: {conf:.2f}"

                        # 计算文本位置
                        text_x = int(box_coords[0])
                        text_y = int(box_coords[1] - 5)  # 在框上方显示文本

                        # 绘制文本
                        cv2.putText(image, title,
                                  (text_x, text_y),
                                  cv2.FONT_HERSHEY_SIMPLEX,
                                  0.5, (0, 255, 0), 2)
    return detectVals, image



def predict_single_img_file(fn, conf_threshold=0.7, plot_result=False):
    """
       根据文件名进行单张图片推理分割
       :param fn: 图片文件路径
       :param conf_threshold: 置信度阈值
       :param plot_result: 是否绘制结果
       :return: 检测结果和处理后的图片
   """
    # 读取图片文件并转换为RGB格式
    image = np.array(Image.open(fn).convert("RGB"))

    # 调用 predict_single_img 进行处理
    return predict_single_img(image, conf_threshold, plot_result)


def predict_single_img(image, conf_threshold=0.7, plot_result=False):
    """
    单张图片推理分割
    :param plot_result: 是否输出推理图片，默认False
    :param image: 需要进行推理的原图
    :param conf_threshold: 置信度
    :return: tuple, cls类别, conf置信度, box外界坐标
    """

    # 获取模型和设备
    model, device = get_model_and_device()
    # 推理
    results = model(image, device=device)
    # 处理检测结果
    return process_detections(results, image, conf_threshold, plot_result)

